#include "ns3/core-module.h"
#include "ns3/network-module.h"
#include "ns3/mobility-module.h"
#include "ns3/internet-module.h"
#include "ns3/applications-module.h"
#include "ns3/wifi-module.h"
#include "ns3/propagation-loss-model.h"
#include "ns3/flow-monitor-module.h"
#include "ns3/config-store-module.h"
#include "ns3/command-line.h"

#include "node-manager.h"
#include "channel-allocation.h"
#include "data-collector.h"
#include "greedy-allocation.h"
#include "random-allocation.h"
#include "epsilon-greedy-allocation.h" 
#include "data-collection-allocation.h" // 确保包含头文件

#include <iostream>
#include <fstream>
#include <string>
#include <vector>

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("SAGINIoVChannelAllocationSim"); // 修改日志组件名称，避免重复

int main(int argc, char *argv[])
{
    // 启用更多的日志组件
    LogComponentEnable("SAGINIoVChannelAllocationSim", LOG_LEVEL_INFO);
    LogComponentEnable("SAGINIoVDataCollector", LOG_LEVEL_INFO);
    LogComponentEnable("NodeManager", LOG_LEVEL_INFO);
    LogComponentEnable("ChannelAllocation", LOG_LEVEL_INFO);
    LogComponentEnable("RandomAllocation", LOG_LEVEL_INFO);
    LogComponentEnable("GreedyAllocation", LOG_LEVEL_INFO);
    
    // 命令行参数
    CommandLine cmd;
    // 修改数据路径为绝对路径
    std::string basePath = "/home/xixi/SAGIN/ns-allinone-3.42/ns-3.42/scratch/SimChannelAlloc-SAGINIoV";
    std::string vehicleDataPath = basePath + "/datas/sumo/preprocessed/output/pingxiang/positions/vehicles";
    std::string satellitesPath = basePath + "/datas/sumo/preprocessed/output/pingxiang/positions/satellites.csv";
    std::string groundStationsPath = basePath + "/datas/sumo/preprocessed/output/pingxiang/positions/ground_stations.csv";
    std::string algorithmType = "random"; // 默认使用随机分配算法
    double simTime = 3600.0; // 仿真时间设为与SUMO数据相同 (0-3600秒)
    double startTime = 0.0;  // 开始时间
    int numChannels = 15; // 可用信道数（10个地面基站+5个卫星基站）
    bool allowSatelliteConnection = true; // 允许车辆连接卫星
    bool allowGroundConnection = true; // 允许车辆连接地面基站
    // 输出目录也使用绝对路径
    std::string outputDir = basePath + "/results/vehicle_traces"; 
    
    // 添加日志显示频率控制
    int logFrequency = 1; // 默认每1秒显示一次日志
    bool enableDetailedLogging = true; // 启用详细日志
    
    // 添加epsilon-greedy参数
    double epsilon = 0.3; // 默认的epsilon值
    
    // 添加数据收集算法参数
    std::string dataCollectionDir = basePath + "/results/data-collection";
    bool useGreedyForFinalDecision = true; // 使用贪心方法作为最终分配决策
    
    cmd.AddValue("vehicleData", "Path to vehicle data directory", vehicleDataPath);
    cmd.AddValue("satellites", "Path to satellites position CSV file", satellitesPath);
    cmd.AddValue("groundStations", "Path to ground stations CSV file", groundStationsPath);
    cmd.AddValue("algorithm", "Channel allocation algorithm (random, greedy, epsilon-greedy, data-collection)", algorithmType);
    cmd.AddValue("simTime", "Simulation time in seconds", simTime);
    cmd.AddValue("startTime", "Simulation start time in seconds", startTime);
    cmd.AddValue("numChannels", "Number of available channels", numChannels);
    cmd.AddValue("allowSatellite", "Allow vehicles to connect to satellites", allowSatelliteConnection);
    cmd.AddValue("allowGround", "Allow vehicles to connect to ground stations", allowGroundConnection);
    cmd.AddValue("outputDir", "Output directory for vehicle traces", outputDir);
    cmd.AddValue("logFreq", "Log display frequency in seconds", logFrequency);
    cmd.AddValue("detailedLog", "Enable detailed vehicle allocation logging", enableDetailedLogging);
    cmd.AddValue("epsilon", "Exploration probability for epsilon-greedy algorithm [0-1]", epsilon);
    cmd.AddValue("dataCollectionDir", "Output directory for data collection", dataCollectionDir);
    cmd.AddValue("useGreedyFinal", "Use greedy method for final allocation decision in data collection", useGreedyForFinalDecision);
    cmd.Parse(argc, argv);
    
    // 确保输出目录存在
    NS_LOG_INFO("Creating output directory if it doesn't exist: " << outputDir);
    std::string createDirCmd = "mkdir -p " + outputDir;
    int dirResult = system(createDirCmd.c_str());
    if (dirResult != 0) {
        NS_LOG_ERROR("Failed to create output directory: " << outputDir);
        return 1;
    }
    
    // 验证目录是否可写
    std::string testFile = outputDir + "/test_write.tmp";
    std::ofstream testStream(testFile.c_str());
    if (!testStream) {
        NS_LOG_ERROR("Output directory is not writable: " << outputDir);
        return 1;
    }
    testStream << "Test" << std::endl;
    testStream.close();
    remove(testFile.c_str());
    
    NS_LOG_INFO("Output directory is ready: " << outputDir);
    
    // 创建节点管理器
    NS_LOG_INFO("Creating node manager");
    NodeManager nodeManager;
    
    // 设置传播模型参数
    // nodeManager.SetGroundStationCoverageRadius(groundStationCoverageRadius);
    // nodeManager.SetSatelliteMaxElevationAngle(satelliteMaxElevationAngle);
    
    // 加载节点数据 - 添加更详细的日志
    NS_LOG_INFO("Loading vehicle data from: " << vehicleDataPath);
    nodeManager.LoadVehicles(vehicleDataPath);
    
    NS_LOG_INFO("Loading satellite data from: " << satellitesPath);
    nodeManager.LoadSatellites(satellitesPath);
    
    NS_LOG_INFO("Loading ground station data from: " << groundStationsPath);
    nodeManager.LoadGroundStations(groundStationsPath);
    
    NS_LOG_INFO("Creating network nodes");
    nodeManager.CreateNodes();
    
    // 获取创建的节点信息并输出
    NodeContainer vehicles = nodeManager.GetVehicleNodes();
    NodeContainer satellites = nodeManager.GetSatelliteNodes();
    NodeContainer groundStations = nodeManager.GetGroundStationNodes();
    
    NS_LOG_INFO("Created " << vehicles.GetN() << " vehicle nodes");
    NS_LOG_INFO("Created " << satellites.GetN() << " satellite nodes");
    NS_LOG_INFO("Created " << groundStations.GetN() << " ground station nodes");
    
    // 创建数据收集器
    NS_LOG_INFO("Setting up data collector");
    DataCollector dataCollector;
    
    // 确保在此处不要重新定义DataCollector的日志组件
    
    dataCollector.EnableThroughputCollection();
    dataCollector.EnableDelayCollection();
    dataCollector.EnableInterferenceCollection();
    dataCollector.EnablePathLossCollection();
    dataCollector.SetSimulationTimeStep(1.0); // 设置为1秒步长，与SUMO数据匹配
    dataCollector.SetOutputDirectory(outputDir);
    
    // 创建信道分配算法
    NS_LOG_INFO("Setting up channel allocation algorithm");
    Ptr<ChannelAllocation> channelAllocator;
    
    if (algorithmType == "random") {
        channelAllocator = CreateObject<RandomAllocation>();
        NS_LOG_INFO("Using Random Allocation algorithm");
    }
    else if (algorithmType == "greedy") {
        channelAllocator = CreateObject<GreedyAllocation>();
        NS_LOG_INFO("Using Greedy Allocation algorithm");
    }
    else if (algorithmType == "epsilon-greedy") {
        // 创建并配置epsilon-greedy分配器
        Ptr<EpsilonGreedyAllocation> epsilonGreedyAllocator = CreateObject<EpsilonGreedyAllocation>();
        epsilonGreedyAllocator->SetEpsilon(epsilon);
        channelAllocator = epsilonGreedyAllocator;
        NS_LOG_INFO("Using Epsilon-Greedy Allocation algorithm with epsilon=" << epsilon);
    }
    else if (algorithmType == "data-collection") {
        // 创建并配置数据收集算法
        Ptr<DataCollectionAllocation> dataCollectionAllocator = CreateObject<DataCollectionAllocation>();
        dataCollectionAllocator->SetOutputDirectory(dataCollectionDir);
        // dataCollectionAllocator->SetUseGreedyForFinalDecision(useGreedyForFinalDecision); // 找到这行代码并注释或删除它
        // 设置NodeManager是在SetDataCollector和SetNodeManager函数中完成的，不需要再显式设置
        channelAllocator = dataCollectionAllocator;
        NS_LOG_INFO("Using Data Collection algorithm");
        
        // 确保输出目录存在
        std::string createDirCmd = "mkdir -p " + dataCollectionDir;
        int dirResult = system(createDirCmd.c_str());
        if (dirResult != 0) {
            NS_LOG_ERROR("Failed to create data collection output directory: " << dataCollectionDir);
            return 1;
        }
        NS_LOG_INFO("Data collection results will be saved to: " << dataCollectionDir);
    }
    else {
        NS_LOG_ERROR("Unknown algorithm type: " << algorithmType);
        return 1;
    }
    
    // 设置算法参数
    channelAllocator->SetNumberOfChannels(numChannels);
    channelAllocator->SetNodeManager(&nodeManager);  // 这行代码已经存在，确保它正确设置
    channelAllocator->SetDataCollector(&dataCollector);
    channelAllocator->SetAllowSatelliteConnection(allowSatelliteConnection);
    channelAllocator->SetAllowGroundStationConnection(allowGroundConnection);
    
    // 准备仿真 - 添加更多日志
    NS_LOG_INFO("Installing internet stack on nodes");
    nodeManager.InstallInternetStack();
    NS_LOG_INFO("Installing applications on nodes");
    nodeManager.InstallApplications();
    
    // 启动周期性数据收集
    NS_LOG_INFO("Setting up periodic data collection");
    dataCollector.StartPeriodicCollection(startTime, simTime);
    dataCollector.EnableRealTimeLogging(enableDetailedLogging, logFrequency); // 启用实时日志
    
    // 启动信道分配
    NS_LOG_INFO("Starting channel allocation");
    channelAllocator->StartPeriodic(Seconds(1.0)); // 每1秒执行一次信道分配
    channelAllocator->EnableRealTimeLogging(enableDetailedLogging, logFrequency); // 启用实时日志
    
    // 运行仿真
    NS_LOG_INFO("Starting simulation for " << simTime << " seconds");
    Simulator::Run();
    
    // 保存结果
    NS_LOG_INFO("Saving results to: " << outputDir);
    std::string resultFile = outputDir + "/results_" + algorithmType + "_aggregate.csv";
    dataCollector.SaveResults(resultFile);
    NS_LOG_INFO("Saved aggregate results to: " << resultFile);
    
    dataCollector.SaveVehicleResults(); // 保存按车辆分组的结果
    NS_LOG_INFO("Saved vehicle results to directory: " << outputDir);
    
    // 清理
    Simulator::Destroy();
    
    NS_LOG_INFO("Simulation completed");
    return 0;
}
